import numpy

def ideal_gas_w2q(gamma, w):
    """Convert the primitive variables `W` to the conserved variables
    `Q` using the ideal gas equation of state.

    Parameters:
    -----------

    gamma : float
        The adiabatic constant for the gas.

    w : numpy.ndarray shape = (3, 1)
        Vector of primitive variables (rho, u, p)

    Returns : numpy.ndarray shape = (3,1) `q` of conserved variables
    (rho, M, E)

    """
    q = numpy.zeros_like(w)
    rho = w[0]; u = w[1]; p = w[2]

    q[0] = rho
    q[1] = rho * u
    q[2] = rho * 0.5 * u * u + p/(gamma - 1.0)

    return q

def ideal_gas_q2w(gamma, q):
    """Convert the conserved variables `Q` to the primitive variables
    `W` using the ideal gas equation of state.

    """
    w = numpy.zeros_like(q)
    w[0] = rho = q[0]; w[1] = u = q[1]/rho; e = q[2]/rho - 0.5 * u * u
    w[2] = p = e * rho * (gamma - 1.0)

    return w

def ideal_gas_flux_q(gamma, q):
    """Compute the flux for the Euler equations using the ideal gas
    equation, from the vector of conserved variables.

    Parameters:
    -----------

    gamma : float
        Adiabatic constant for the gas.

    q : array
        Vector of conserved variables.

    Returns : array of flux values

    """
    f = numpy.zeros_like(q)

    rho = q[0]
    u = q[1]/rho
    e = q[2]/rho - 0.5 * u * u; p = (gamma - 1.0) * e * rho

    f[0] = q[1]
    f[1] = q[1] * u + p
    f[2] = u * (q[2] + p)

    return f
